import akshare as ak
import pandas as pd
import os
import argparse
from datetime import datetime, timedelta
import time

# 配置请求头（解决反爬）[4](@ref)
ak.tool.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"

def get_target_stock_data(stock_id: str, days: int, period: str) -> pd.DataFrame:
    """获取指定证券最近N天数据"""
    try:
        # 时间范围扩展为自然日（含非交易日）[6](@ref)
        start_date = (datetime.now() - timedelta(days=days*2)).strftime("%Y-%m-%d %H:%M:%S")
        end_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # 统一周期参数处理（去除min后缀）[6](@ref)
        period_num = period.replace("min", "") 

        # ETF处理（使用通用分钟级接口）
        if stock_id.startswith(('51', '15')): 
            df = ak.fund_etf_hist_min_em(
                symbol=stock_id,
                period=period_num,
                start_date=start_date,
                end_date=end_date,
                adjust="hfq"
            )
            if not df.empty:
                # 字段标准化[1,6](@ref)
                df = df.rename(columns={
                    '时间': 'datetime',
                    '收盘': 'close',
                    '最高': 'high',
                    '最低': 'low',
                    '成交量': 'volume'
                })
        
        # 沪深A股处理（适配新时间格式）
        elif stock_id.startswith(('6','0','3')):
            market = "sh" if stock_id.startswith('6') else "sz"
            # 使用统一分钟接口[5,7](@ref)
            df = ak.stock_zh_a_hist_min_em(
                symbol=f"{stock_id}",
                period=period_num,
                start_date=start_date,
                end_date=end_date,
                adjust="hfq"
            )
            # 字段重命名
            df = df.rename(columns={'时间': 'datetime'})
        
        # 数据有效性过滤（排除非交易时段）
        if not df.empty:
            df['datetime'] = pd.to_datetime(df['datetime'])
            # 按自然日过滤（保留最后N天的数据）[6](@ref)
            cutoff_date = datetime.now() - timedelta(days=days)
            return df[df['datetime'] >= cutoff_date].sort_values('datetime')
        return pd.DataFrame()
    
    except Exception as e:
        print(f"获取 {stock_id} 失败: {str(e)}")
        return pd.DataFrame()

if __name__ == "__main__":
    # 参数说明更新[6](@ref)
    parser = argparse.ArgumentParser(description='证券数据下载工具')
    parser.add_argument('--id', required=True, help='证券代码（如：513260）')
    parser.add_argument('--time', type=int, default=30, 
                      help='获取最近N天数据（默认30天，建议不超过60天）[4](@ref)')
    parser.add_argument('--period', type=str, default="1min", 
                      help='K线周期（可选：1min/5min/15min/30min/60min）')
    args = parser.parse_args()

    os.makedirs('stock_data', exist_ok=True)
    
    # 智能重试机制（增加间隔时间）[4](@ref)
    max_retries = 3
    for attempt in range(max_retries):
        df = get_target_stock_data(args.id, args.time, args.period)
        if not df.empty:
            break
        print(f"第{attempt+1}次重试（建议检查网络或缩短时间范围）...")
        time.sleep(10)
    
    if not df.empty:
        # 文件名格式优化[6](@ref)
        filename = f"stock_data/{args.id}_{args.time}days_{args.period}.csv"
        df.to_csv(filename, index=False)
        print(f"成功保存 {args.id} 最近{args.time}天数据至 {filename}")
        print(f"数据时间范围：{df['datetime'].min()} 至 {df['datetime'].max()}")
        print(f"数据量：{len(df)} 条（{args.period}周期）")
    else:
        print("数据获取失败，可能原因：")
        print("1. 接口限制：部分品种1分钟数据仅支持当日查询（建议缩短时间范围）[6](@ref)")
        print("2. 代码过期：请尝试更新akshare库（pip install akshare --upgrade）[3](@ref)")
        print("3. 网络限制：建议使用代理IP池（参考反爬解决方案）[4](@ref)")